﻿// Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <iostream>
#include <ostream>
#include <vector>

#include "visualize.h"

namespace ai {

std::vector<std::string> ReadDict(const std::string &path) {
  std::ifstream in(path);
  std::string line;
  std::vector<std::string> m_vec;
  if (in) {
    while (getline(in, line)) {
      m_vec.push_back(line);
    }
  } else {
    spdlog::get("logger")->info("no such label file: {}, exit the program...", path);
    exit(1);
  }
  m_vec.insert(m_vec.begin(), "#"); // blank char for ctc
  m_vec.push_back(" ");
  return m_vec;
}


void GetRotateCropImage(const cv::Mat srcimage,
  std::vector<cv::Point> four_points,
  cv::Mat &crop_image) 
{
  float x_collect[4] = {four_points[0].x, four_points[1].x, four_points[2].x, four_points[3].x};
  float y_collect[4] = {four_points[0].y, four_points[1].y, four_points[2].y, four_points[3].y};
  int left = int(*std::min_element(x_collect, x_collect + 4));
  int right = int(*std::max_element(x_collect, x_collect + 4));
  int top = int(*std::min_element(y_collect, y_collect + 4));
  int bottom = int(*std::max_element(y_collect, y_collect + 4));

  if (left < 0) {left = 0;}
  if (top < 0) {top = 0;}
  if (right > srcimage.cols) {right = srcimage.cols;}
  if (bottom > srcimage.rows) {bottom = srcimage.rows;}

  cv::Mat img_crop = srcimage(cv::Rect(left, top, right - left, bottom - top));

  for (int i = 0; i < four_points.size(); i++) {
    four_points[i] = cv::Point(four_points[i].x - left, four_points[i].y - top);
  }

  int img_crop_width = int(sqrt(pow(four_points[0].x - four_points[1].x, 2) +
                                pow(four_points[0].y - four_points[1].y, 2)));
  int img_crop_height = int(sqrt(pow(four_points[0].x - four_points[3].x, 2) +
                                 pow(four_points[0].y - four_points[3].y, 2)));

  cv::Point2f pts_std[4];
  pts_std[0] = cv::Point2f(0., 0.);
  pts_std[1] = cv::Point2f(img_crop_width, 0.);
  pts_std[2] = cv::Point2f(img_crop_width, img_crop_height);
  pts_std[3] = cv::Point2f(0.f, img_crop_height);

  cv::Point2f pointsf[4];
  pointsf[0] = cv::Point2f(four_points[0].x, four_points[0].y);
  pointsf[1] = cv::Point2f(four_points[1].x, four_points[1].y);
  pointsf[2] = cv::Point2f(four_points[2].x, four_points[2].y);
  pointsf[3] = cv::Point2f(four_points[3].x, four_points[3].y);

  cv::Mat M = cv::getPerspectiveTransform(pointsf, pts_std);

  cv::Mat dst_img;
  cv::warpPerspective(img_crop, dst_img, M,
                      cv::Size(img_crop_width, img_crop_height),
                      cv::BORDER_REPLICATE);

  // // 针对行驶证 仅有水平文字
  // if (dst_img.cols * 1.0 / dst_img.rows > 2.0){
  //   crop_image = dst_img;
  // }else{
  //   cv::Mat srcCopy = cv::Mat(dst_img.rows, dst_img.cols, dst_img.depth());
  //   cv::transpose(dst_img, srcCopy);
  //   cv::flip(srcCopy, srcCopy, 0);
  //   crop_image = srcCopy;
  // }
  crop_image = dst_img.clone();

}


std::vector<int> GenerateColorMap(int num_class) {
  auto colormap = std::vector<int>(3 * num_class, 0);
  for (int i = 0; i < num_class; ++i) {
    int j = 0;
    int lab = i;
    while (lab) {
      colormap[i * 3] |= (((lab >> 0) & 1) << (7 - j));
      colormap[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j));
      colormap[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j));
      ++j;
      lab >>= 3;
    }
  }
  std::reverse(colormap.begin(), colormap.end());
  return colormap;
}


cv::Mat visualize(const cv::Mat& img,
                  std::vector<BaseInfo*>& det_infos,
                  const std::vector<std::string>& labels) {
  auto colormap = GenerateColorMap(labels.size());
  cv::Mat vis_img = img.clone();
  for (int i = 0; i < det_infos.size(); ++i) {
    DetInfo* det_info = (DetInfo*)(det_infos[i]);
    // draw box and title
    std::string text = det_info->category;
    int c1 = colormap[3 * det_info->class_idx + 0];
    int c2 = colormap[3 * det_info->class_idx + 1];
    int c3 = colormap[3 * det_info->class_idx + 2];
    cv::Scalar box_color = cv::Scalar(c1, c2, c3);
    text += std::to_string(static_cast<int>(det_info->score * 100)) + "%";
    int font_face = cv::FONT_HERSHEY_SIMPLEX;
    double font_scale = 0.5f;
    float thickness = 0.5;
    cv::Size text_size =
        cv::getTextSize(text, font_face, font_scale, thickness, nullptr);
    cv::Point origin = cv::Point(det_info->det_box.x, det_info->det_box.y);

    // background
    cv::Rect text_back = cv::Rect(det_info->det_box.x,
                                  det_info->det_box.y - text_size.height,
                                  text_size.width,
                                  text_size.height);

    // draw
    cv::rectangle(vis_img, det_info->det_box, box_color, 2);
    cv::rectangle(vis_img, text_back, box_color, -1);
    cv::putText(vis_img,
                text,
                origin,
                font_face,
                font_scale,
                cv::Scalar(255, 255, 255),
                thickness);

    // mask
    if (det_info->mask.empty()) { continue; }
    // std::vector<float> mask_data;
    // mask_data.assign(det_infos[i].mask., det_infos[i].mask.data.end());
    // cv::Mat bin_mask(det_infos[i].mask.shape[1],
    //                  det_infos[i].mask.shape[0],
    //                  CV_32FC1,
    //                  mask_data.data());
    // cv::Mat full_mask = cv::Mat::zeros(vis_img.size(), CV_8UC1);
    // bin_mask.copyTo(full_mask(roi));
    // cv::Mat mask_ch[3];
    // mask_ch[0] = full_mask * c1;
    // mask_ch[1] = full_mask * c2;
    // mask_ch[2] = full_mask * c3;
    // cv::Mat mask;
    // cv::merge(mask_ch, 3, mask);
    // cv::addWeighted(vis_img, 1, mask, 0.5, 0, vis_img);
  }
  return vis_img;
}

void visualize_db( cv::Mat& srcimg,
                      cv::Mat &save_image,
                      std::vector <BaseInfo*> &text_lines,
                      bool save_img) {
  srcimg.copyTo(save_image);
  for (int n = 0; n < text_lines.size(); n++) {
    TextDetInfo* tdi = (TextDetInfo*)(text_lines[n]);
    // std::cout << "visualize_db: " << tdi->four_points[0] << "  " << tdi->four_points[1] << "  " << tdi->four_points[2] << "  " << tdi->four_points[3] << std::endl;

    std::vector<std::vector<cv::Point>> fp = {tdi->four_points};
    cv::drawContours(save_image, fp, 0, cv::Scalar(0,255,0));
  //   cv::Point rook_points[4];
  //   for (int m = 0; m < 4; m++) {
  //     rook_points[m] = cv::Point(int(tdi->four_points[m].x), int(tdi->four_points[m].y));
  //   }

  //   const cv::Point *ppt[1] = {rook_points};
  //   int npt[] = {4};
  //   cv::polylines(save_image, ppt, npt, 1, 1, CV_RGB(0, 255, 0), 1, 8, 0);
  // }

  // if (save_img){
  //   cv::imwrite("./ocr_vis.png", save_image);
  //   spdlog::get("logger")->info("The detection visualized image saved in ./ocr_vis.png");
  }
}

// Visualiztion MaskDetector results
cv::Mat draw_det_box(cv::Mat& img,
                        const std::vector<DetInfo>& results,
                        const std::vector<std::string>& lable_list,
                        const std::vector<int>& colormap) {
  cv::Mat vis_img = img.clone();
  for (int i = 0; i < results.size(); ++i) {

    // Configure color and text size
    std::ostringstream oss;
    oss << std::setiosflags(std::ios::fixed) << std::setprecision(4);
    oss << lable_list[results[i].class_idx] << " ";
    oss << results[i].score;
    std::string text = oss.str();
    int c1 = colormap[3 * results[i].class_idx + 0];
    int c2 = colormap[3 * results[i].class_idx + 1];
    int c3 = colormap[3 * results[i].class_idx + 2];
    cv::Scalar roi_color = cv::Scalar(c1, c2, c3);
    int font_face = cv::FONT_HERSHEY_COMPLEX_SMALL;
    double font_scale = 0.5f;
    float thickness = 0.5;
    cv::Size text_size = cv::getTextSize(text,
                                         font_face,
                                         font_scale,
                                         thickness,
                                         nullptr);
    cv::Point origin = results[i].det_box.tl();

    // Configure text background
    cv::Rect text_back = cv::Rect(results[i].det_box.x,
                                  results[i].det_box.y - text_size.height,
                                  text_size.width,
                                  text_size.height);

    // Draw roi object, text, and background
    cv::rectangle(vis_img, results[i].det_box, roi_color, 2);
    cv::rectangle(vis_img, text_back, roi_color, -1);
    cv::putText(vis_img,
                text,
                origin,
                font_face,
                font_scale,
                cv::Scalar(255, 255, 255),
                thickness);
  }
  return vis_img;
}

void log_tensor_shape(const std::string& name, std::vector<std::string> data) {
    std::string out_str =  name + ": [";
    for (const auto& n : data) { out_str += n + ", "; }
    spdlog::get("logger")->info("{}]", out_str);
}
} // namespace paddle_infer